# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from experimental.prophet.graph import graph
from experimental.prophet.ops.util import (to_tuple,
                                           to_data_type_proto,
                                           to_duration_int,
                                           to_duration_str)
from experimental.prophet.contrib.functional import vectorize


@vectorize
def shift_y(value, window_size, auto_name=True, namespace='Y'):
  assert window_size > 0, window_size
  value = graph.as_variable(value)

  node = graph.create_node('y_gen.Shift', [value],
                           attrs={
                               'dtype': to_data_type_proto(value.dtype), 'window_size': window_size
                           })
  out_var = node.outputs[0]

  if auto_name and value.is_named_ref:
    out_var = out_var.named_ref(value.name + '.Y' + str(window_size) + 'shift', namespace=namespace)

  return out_var


@vectorize
def time_shift_y(value, window_size, timestamp=None, auto_name=True, namespace='Y'):
  if timestamp is None:
    timestamp = graph.default_timestamp()

  value = graph.as_variable(value)
  window_size = to_duration_int(window_size)
  assert window_size >= 1000000, window_size  # 1ms

  node = graph.create_node('y_gen.TimeShift', [timestamp, value],
                           attrs={
                               'dtype': to_data_type_proto(value.dtype),
                               'window_size_ns': window_size
                           })
  out_var = node.outputs[0]

  if auto_name and value.is_named_ref:
    out_var = out_var.named_ref(value.name + '.Y' + to_duration_str(window_size),
                                namespace=namespace)

  return out_var
